# 导入必备的工具包
import torch

# 预定义的网络层torch.nn, 工具开发者已经帮助我们开发好的一些常用层,
# 比如，卷积层, lstm层, embedding层等, 不需要我们再重新造轮子.
import torch.nn as nn

# 数学计算工具包
import math

# torch中变量封装函数Variable.
from torch.autograd import Variable

import warnings
warnings.filterwarnings('ignore')
# 文本嵌入层的作用
# 无论是源文本嵌入还是目标文本嵌入，都是为了将文本中词汇的数字表示转变为向量表示, 希望在这样的高维空间捕捉词汇间的关系.

# 定义Embeddings类来实现文本嵌入层，这里s说明代表两个一模一样的嵌入层, 他们共享参数.
# 该类继承nn.Module, 这样就有标准层的一些功能, 这里我们也可以理解为一种模式, 我们自己实现的所有层都会这样去写.
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        """类的初始化函数, 有两个参数, d_model: 指词嵌入的维度, vocab: 指词表的大小."""
        # 接着就是使用super的方式指明继承nn.Module的初始化函数, 我们自己实现的所有层都会这样去写.
        super(Embeddings, self).__init__()
        # 之后就是调用nn中的预定义层Embedding, 获得一个词嵌入对象self.lut
        self.lut = nn.Embedding(vocab, d_model)
        # 最后就是将d_model传入类中
        self.d_model = d_model

    def forward(self, x):
        """可以将其理解为该层的前向传播逻辑，所有层中都会有此函数
           当传给该类的实例化对象参数时, 自动调用该类函数
           参数x: 因为Embedding层是首层, 所以代表输入给模型的文本通过词汇映射后的张量"""

        # 将x传给self.lut并与根号下self.d_model相乘作为结果返回
        return self.lut(x) * math.sqrt(self.d_model)

'''
# 创建一个包含10个词的词嵌入(词汇表的长度)，每个词嵌入的维度为5（每个词的维度）
embedding_layer = nn.Embedding(num_embeddings=10, embedding_dim=5)
# 结果
tensor([[-0.4150, -0.7601,  0.2033,  0.5462,  0.3505],
        [ 1.1467, -0.5948, -0.7772, -0.8821, -0.0594],
        [ 0.5237, -0.3369,  1.6435,  0.2094, -0.3540],
        [ 0.7534,  2.1488, -1.4210, -0.2166,  2.2788],
        [ 0.1599,  0.2382,  1.0079, -1.4361, -0.1213],
        [ 0.1802, -0.3412, -1.0486, -0.8906, -1.0458],
        [-2.0671,  1.5961,  0.4305, -0.3616, -1.1633],
        [-0.9917, -0.4111,  0.5032,  1.3175,  0.1757],
        [ 1.4751, -0.1087,  2.9269,  1.2613,  1.6095],
        [-1.1523,  0.9529, -1.2208,  1.3535,  1.3478]]

d_model=10表示词嵌入维度，当选取的维度过低时， 词嵌入 将有较大的偏差（high bias），因为模型丢失了较多信号。而维度过高时， 词嵌入的效果将会变差，因为模型会过度拟合训练数据，无法泛化到新的数据上
embedding = nn.Embedding(10, 3)
input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])

每个数字会被映射成3个维度，[-0.1267,  2.2847, -0.6410]代表1， [ 0.2273,  0.0161, -0.8555]代表2
tensor([[[-0.1267,  2.2847, -0.6410],
         [ 0.2273,  0.0161, -0.8555],
         [-0.6550, -1.1668, -1.4832],
         [ 0.1156,  0.7057,  1.0474]],
        [[-0.6550, -1.1668, -1.4832],
         [ 0.4194, -2.7155,  0.0739],
         [ 0.2273,  0.0161, -0.8555],
         [-0.1732,  0.0813, -0.1582]]])
         
         
'''

embedding = nn.Embedding(10, 3)
input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
embedding1 = embedding(input)
print(embedding1)

embedding = nn.Embedding(10, 3, padding_idx=0)
input = torch.LongTensor([[0, 2, 0, 5]])
embedding2 = embedding(input)
print(embedding2.shape)

# 词嵌入维度是512维
d_model = 512

# 词表大小是1000
vocab = 1000
# 输入x是一个使用Variable封装的长整型张量, 形状是2 x 4
x = Variable(torch.LongTensor([[100, 2, 421, 508], [491, 998, 1, 221]]))
emb = Embeddings(d_model, vocab)
embr = emb(x)
print("embr:", embr.shape)

if __name__ == '__main__':
    print("over")
